package geppetto.phraseTable.phrase.feature.calc.local;

import geppetto.cat.corpus.BilingualCorpus;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.ExtractedPhrasePairData;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.data.ReorderingData;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.data.ReorderingDataCalc;
import geppetto.phraseTable.PhraseTable;
import geppetto.phraseTable.phrase.Phrase;
import geppetto.phraseTable.phrase.feature.AbstractFeature;
import geppetto.phraseTable.phrase.feature.PenaltyFeature;
import geppetto.phraseTable.phrase.feature.ReorderingFeature;
import geppetto.reordering.MSDReorderingCalculator;
import geppetto.reordering.ReorderingInfo;
import gnu.trove.TIntDoubleHashMap;
import gnu.trove.TIntObjectHashMap;

import java.io.PrintStream;
import java.util.ArrayList;
import java.util.HashMap;


public class ReorderingFeatureCalc implements AbstractFeatureCalc{
	protected final static String REORDERING_HEADER = "reordering";
	protected final static String SCORE_HEADER = "score";
	protected final static String FEATURE = "reord";

	protected static int MONO_CODE = 1;
	protected static int SWAP_CODE = 2;
	protected static int OTHER_CODE = 3;
	protected static int DRIGHT_CODE = 4;
	protected static int DLEFT_CODE = 5;
	protected static int NOMONO_CODE = 6;
	protected TIntDoubleHashMap prevSmoothingFactors = new TIntDoubleHashMap();
	protected TIntDoubleHashMap nextSmoothingFactors = new TIntDoubleHashMap();
	protected boolean shouldPrint = true;
	protected int order;
		
	protected static HashMap<Integer, String> intToStringHash = new HashMap<Integer, String>();
	protected static double DEFAULTSMOOTHING = 0.5;
	
	public ReorderingFeatureCalc(int order) {
		this(order, DEFAULTSMOOTHING);
	}
	
	public ReorderingFeatureCalc(int order, double smoothing) {
		super();
		setSmoothing(smoothing);
		this.order = order;
	}
	
	public ReorderingFeatureCalc(int order, double smoothing, boolean shouldPrint) {
		super();
		setSmoothing(smoothing);
		this.shouldPrint = shouldPrint;
		this.order = order;
	}

	public void setSmoothing(double smoothing){
		prevSmoothingFactors.put(MONO_CODE, smoothing);
		prevSmoothingFactors.put(SWAP_CODE, smoothing);
		prevSmoothingFactors.put(OTHER_CODE, smoothing);
		prevSmoothingFactors.put(DRIGHT_CODE, smoothing);
		prevSmoothingFactors.put(DLEFT_CODE, smoothing);
		prevSmoothingFactors.put(NOMONO_CODE, smoothing);
		nextSmoothingFactors.put(MONO_CODE, smoothing);
		nextSmoothingFactors.put(SWAP_CODE, smoothing);
		nextSmoothingFactors.put(OTHER_CODE, smoothing);
		nextSmoothingFactors.put(DRIGHT_CODE, smoothing);
		nextSmoothingFactors.put(DLEFT_CODE, smoothing);
		nextSmoothingFactors.put(NOMONO_CODE, smoothing);
	}
	
	public void printMosesReorderingTable(PrintStream out, BilingualCorpus corpus, byte sentenceSource, PhraseTable pt, int maxTargetDuration, int maxSourceDuration){		
		for(int targetLen = 0; targetLen < maxTargetDuration; targetLen++){
			for(int sourceLen = 0; sourceLen < maxSourceDuration; sourceLen++){
				TIntObjectHashMap<TIntObjectHashMap<Phrase>> phrasesT = pt._phraseTable[sourceLen][targetLen];
				int keysT[] = phrasesT.keys();
				for (int i = 0; i < keysT.length; i++) {
					TIntObjectHashMap<Phrase> phrases = phrasesT.get(keysT[i]);
					int[] keys = phrases.keys();
					for (int j = 0; j < keys.length; j++) {
						Phrase p = phrases.get(keys[j]);
						double[] reordPrevValues = getPhrasePrevReorderingFeatures(p.get_reorderingInfo());
						double[] reordNextValues = getPhraseNextReoderingFeatures(p.get_reorderingInfo());						
						
						out.println(p.stringMosesReordering(corpus, pt._sourcePhrasesVocab, pt._targetPhrasesVocab, reordPrevValues, reordNextValues));
					}
				}
			}
		}
	}
	
	public double[] getPhrasePrevReorderingFeatures(ReorderingInfo info){
		int[] params = new int[]{MONO_CODE, SWAP_CODE, OTHER_CODE};
		double[] reordPrevValues = new double[params.length];
		
		double total = 0;
		ReorderingInfo.normalize(info);
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = prevSmoothingFactors.get(param);
			double reordValue = info.getPreviousWordOrientationCounts().get(param);
			reordPrevValues[k] = smoothing + reordValue;
			total += reordPrevValues[k];
		}						
		for (int k = 0; k < params.length; k++){
			reordPrevValues[k] = reordPrevValues[k] / total;
		}
		
		return reordPrevValues;
	}
	
	public double[] getPhraseNextReoderingFeatures(ReorderingInfo info){
		int[] params = new int[]{MONO_CODE, SWAP_CODE, OTHER_CODE};
		double[] reordNextValues = new double[params.length];

		double total = 0;
		ReorderingInfo.normalize(info);
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = nextSmoothingFactors.get(param);
			double reordValue = info.getNextWordOrientationCounts().get(param);
			reordNextValues[k] = smoothing + reordValue;
			total += reordNextValues[k];
		}
		for (int k = 0; k < params.length; k++){
			reordNextValues[k] = reordNextValues[k] / total;
		}
		return reordNextValues;
	}
	
	public double[] getPhraseReorderingFeatures(ReorderingInfo info){
		int[] params = new int[]{MONO_CODE, SWAP_CODE, OTHER_CODE};
		double[] reordValues = new double[params.length*2];
		
		double total = 0;
		//info.normalize();
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = prevSmoothingFactors.get(param);
			double reordValue = info.getPreviousWordOrientationCounts().get(param);
			reordValues[k] = smoothing + reordValue;
			total += reordValues[k];
		}						
		for (int k = 0; k < params.length; k++){
			reordValues[k] = reordValues[k] / total;
		}
		
		total = 0;
		for (int k = 0; k < params.length; k++){
			int param = params[k];
			double smoothing = nextSmoothingFactors.get(param);
			double reordValue = info.getNextWordOrientationCounts().get(param);
			reordValues[k+params.length] = smoothing + reordValue;
			total += reordValues[k+params.length];
		}
		for (int k = 0; k < params.length; k++){
			reordValues[k+params.length] = reordValues[k+params.length] / total;
		}
		
		return reordValues;
	}

	@Override
	public void calculateFeature(ArrayList<Phrase> phrases) {
		for(int i = 0; i < phrases.size(); i++){
			Phrase p = phrases.get(i);
			ReorderingData rdata = (ReorderingData)p.getData(REORDERING_HEADER);
			phrases.get(i)._features.addFeature(FEATURE, new ReorderingFeature(getPhraseReorderingFeatures(rdata.getReorderingInfo())),order,shouldPrint);
		}
	}
}
